# -*- coding: UTF-8 -*-
import os
import threading
import yaml
import numpy as np
import rospy
import cv2
from sensor_msgs.msg import Image
from sensor_msgs.msg import Imu
from cv_bridge import CvBridge
import module.IMU_Optimize
import module.lowlight_enhancement
import queue
import time

#pip install opencv-python==3.4.3.18

bridge = CvBridge()
imuMQ=queue.Queue()
imageMQ=queue.Queue()
imuList=[]
imageList=[]
imuCheckingList=[]
imgProcessingList=[]

imuSendList=[]
imageSendList=[]

imu_buf=threading.Lock()
img_buf=threading.Lock()
processing=threading.Lock()
processingImage=threading.Lock()

data_train=None
imu_topic=None
imu_pub=None
image_topic=None
image_pub=None

last_imu_template=[]


def initParameter():
    global  imu_topic,image_topic,imu_pub,image_pub
    fs = open(os.path.join("./config", "parameter.yaml"), encoding="UTF-8")
    datas = yaml.load(fs, Loader=yaml.FullLoader)
    imu_topic=datas['imu_topic']
    image_topic=datas["image_topic"]
    imu_pub=datas["imu_pub"]
    image_pub=datas["image_pub"]


def handleIMU(data):
    imu_buf.acquire()
    imuMQ.put(data)
    imu_buf.release()

def handleImage(data):
    img_buf.acquire()
    imageMQ.put(data)
    img_buf.release()

def imgProcess():
    global imageList,imgProcessingList
    imgProcessingList=imageList
    imageList=[]
    template=[]
    for data in imgProcessingList:
        cv_image = bridge.imgmsg_to_cv2(data, desired_encoding='bgr8')
        from PIL import Image
        frame = Image.fromarray(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
        result=module.lowlight_enhancement.lowlight(frame)
        img = cv2.cvtColor(np.asarray(result), cv2.COLOR_RGB2BGR)
        image_message = bridge.cv2_to_imgmsg(img, encoding="bgr8")
        template.append(image_message)
    for frame,target in zip(imgProcessingList,template):
        frame.data=target.data
        frame.encoding=target.encoding
        frame.height=target.height
        frame.width=target.width
        frame.step=target.step


def imuProcessByDTW():
    global imuList,imuCheckingList,last_imu_template
    imuCheckingList=imuList
    imuList=[]
    scores=-1
    if len(last_imu_template)==0:
        last_imu_template=imuCheckingList.copy()
    else:
        scores=module.IMU_Optimize.estimate_dtw(imuCheckingList,last_imu_template)
        last_imu_template=imuCheckingList.copy()
        if scores>=2.0:
           imuCheckingList=module.IMU_Optimize.data_replacing(imuCheckingList)
    print("DTW done! score : "+str(scores))
    #print("DTW done!")



def publishImuByGroup():
    global imuSendList, imuCheckingList
    imuSendList=imuCheckingList
    imuCheckingList=[]
    imuSendList.reverse()
    while len(imuSendList)>0:
        imupub.publish(imuSendList.pop())
        imurate.sleep()

def publishImgByGroup():
    global imgProcessingList,imageSendList
    imageSendList=imgProcessingList
    imgProcessingList=[]
    imageSendList.reverse()
    while len(imageSendList)>0:
        imgpub.publish(imageSendList.pop())
        imgrate.sleep()

def dataTransfer(maxIMU,maxImage):
    global imuList,imuMQ,imageList,imageMQ
    for i in range(0,maxIMU):
        imuList.append(imuMQ.get())
    for j in range(0,maxImage):
        imageList.append(imageMQ.get())

def checkPublishStatus():
    global i
    lastPublish=0
    while True:
        current=time.perf_counter()
        if imuMQ.qsize()>=300 and imageMQ.qsize() >= 30:
                imu_buf.acquire()
                img_buf.acquire()
                dataTransfer(300,30)
                imu_buf.release()
                img_buf.release()
        # elif current-lastPublish> 3:
        #     print("long time no publish")
        #     continue
        else: continue
        imuCheckTask = threading.Thread(target=imuProcessByDTW)
        imgProcessTask = threading.Thread(target=imgProcess)
        imuCheckTask.start()
        imgProcessTask.start()
        imuCheckTask.join()
        imgProcessTask.join()
        print("imu & image processed")
        imuPubTask = threading.Thread(target=publishImuByGroup)
        imgPubTask = threading.Thread(target=publishImgByGroup)
        imuPubTask.start()
        imgPubTask.start()
        imuPubTask.join()
        imgPubTask.join()
        lastPublish=current
        print("imu & img publish done")
        print("--------------------------------------")


if __name__ == "__main__":

    rospy.init_node('FeatureProject',anonymous=True)
    initParameter()
    imusub = rospy.Subscriber(imu_topic, Imu,handleIMU, queue_size=20)
    imgsub=rospy.Subscriber(image_topic,Image,handleImage,queue_size=20)
    imupub=rospy.Publisher(imu_pub,Imu,queue_size=20)
    imgpub=rospy.Publisher(image_pub,Image,queue_size=20)

    imurate = rospy.Rate(300)
    imgrate = rospy.Rate(30)

    #
    # files = get_imu_record("imu_template")
    # for file in files:
    #     data_train = np.load(file)

    featureProject=threading.Thread(target=checkPublishStatus)
    featureProject.start()
    rospy.spin()

